#!/usr/bin/env python

from Bio import SeqIO

import argparse

import sys

dis = '''
从 gtf文件中 提取 3'UTR
用法 ./gtf23utr.py -g test.gtf -f test.fa -o utr
'''


parser = argparse.ArgumentParser(description=dis)

parser.add_argument('-g', type=str, nargs='+',
    help='gtf') # 这里metavar指赋值的参数的说明 nargs 指只能给定一个

parser.add_argument('-f', type=str, nargs='+',
    help='fasta文件') 

parser.add_argument('-o', type=str, nargs='+',
    help='输出文件的前缀') 



# 获取参数
args = parser.parse_args()

if args.f == None:
    print('请给定fasta文件')
    parser.print_help()
    sys.exit(888)
else:
    infile_fasta = args.f[0]

if args.g == None:
    print('请给定gtf文件')
    parser.print_help()
    sys.exit(888)
else:
    infile_gtf = args.g[0]

if args.o == None:
    print('请给定输出文件')
    parser.print_help()
    sys.exit(888)
else:
    outfile = args.o[0]



class gtf_unit(object):

    def __init__(self, gff_lista):

        self.scaf = gff_lista[0]

        self.sorce = gff_lista[1]
        
        self.type = gff_lista[2]
        
        self.start = int(gff_lista[3])-1
        
        self.end = int(gff_lista[4])

        self.strand = gff_lista[6]

        self.frame = gff_lista[7]

        last_dic = {}

        for i in gff_lista[8].strip().split(';'):

            t =  i.split( '"') 

            if len(t)>1:
                last_dic[t[0].strip(' "')] = t[1].strip(' "')


        self.transcript_id = last_dic['transcript_id']

        self.gene_id = last_dic['gene_id'] if 'gene_id' in last_dic else ''
            
        self.gene_name = last_dic['gene_name'] if 'gene_name' in last_dic else ''

    def extract_seq(self,fasta_dic):

        full_seq = fasta_dic[self.scaf]

        self.seq = full_seq[self.start:self.end]

        if (self.strand=='-'):

            self.seq = self.seq.reverse_complement()
        
        

class gtf_gene(object):

    def __init__(self, gtf_unit):
        
        self.gene_id = gtf_unit.gene_id
        
        self.gene_name = gtf_unit.gene_name

        self.transcript = {}


class gtf_trans(object):

    def __init__(self, gtf_unit): #transcript_id, gene_id, gene_name=None):

        self.transcript_id = gtf_unit.transcript_id

        self.gene_id = gtf_unit.gene_id

        self.gene_name = gtf_unit.gene_name

        self.strand = gtf_unit.strand

        self.exons = []

        self.cds = []

        self.mRNA = ''

        self.code = ''

    def fetchSeq(self, fasta_dic):

        if self.strand=='-':
            self.cds = sorted(self.cds, key=lambda x:-x.end)
        else:
            self.cds = sorted(self.cds, key=lambda x:x.start)


        if self.strand=='-':
            self.exons = sorted(self.exons, key=lambda x:-x.end)
        else:
            self.exons = sorted(self.exons, key=lambda x:x.start)


        for i in self.cds:

            i.extract_seq(fasta_dic)

            self.code += i.seq

        for i in self.exons:
            i.extract_seq(fasta_dic)

            self.mRNA += i.seq

        self.utr_3p = ''

        str_code = str(self.code)
        str_mRNA = str(self.mRNA)

        if str_code !='' and str_mRNA.find(str_code)!=-1:

            self.utr_3p = self.mRNA[str_mRNA.find(str_code)+len(self.code):]



    def translation(self):
        self.prot = ''
        if (len(self.code)>=3):
            self.prot = self.code.translate()

    


# 读取基因组文件

fasta_dic = {}

for i in SeqIO.parse(infile_fasta,'fasta'):

    fasta_dic[str(i.id)] = i.seq





gtf_lista = []

with open(infile_gtf) as fila:

    for i in fila:

        k = i.strip().split('\t')


        if len(k)==9 and i[0]!='#':

            gtf_lista.append(gtf_unit(k))


# # first import all gene
# gene_dic = {}

# for i in gtf_lista:

#     if i.type=='gene':

#         gene_dic[i.gene_id] = i


# second import all trans





trans_dic = {}

for i in gtf_lista:

    if i.type.lower() =='transcript':

        trans_dic[i.transcript_id] = gtf_trans(i)

    elif i.type.lower() =='cds':

        trans_dic[i.transcript_id].cds.append(i)
        
    elif i.type.lower() == 'exon':

        trans_dic[i.transcript_id].exons.append(i)


for i in trans_dic:

    trans_dic[i].fetchSeq(fasta_dic)
    trans_dic[i].translation()
    

with open(outfile+'.prot.fa','w') as fila:
    for i in trans_dic:
        item = trans_dic[i]
        fila.write('>'+item.gene_id+':'+item.transcript_id+':prot'+'\n')
        fila.write(str(item.prot)+'\n')

with open(outfile+'.utr.fa','w') as fila:
    for i in trans_dic:
        item = trans_dic[i]
        fila.write('>'+item.gene_id+':'+item.transcript_id+':utr'+'\n')
        fila.write(str(item.utr_3p)+'\n')


with open(outfile+'.cds.fa','w') as fila:
    for i in trans_dic:
        item = trans_dic[i]
        fila.write('>'+item.gene_id+':'+item.transcript_id+':cds'+'\n')
        fila.write(str(item.code)+'\n')

with open(outfile+'.mRNA.fa','w') as fila:
    for i in trans_dic:
        item = trans_dic[i]
        fila.write('>'+item.gene_id+':'+item.transcript_id+':mRNA'+'\n')
        fila.write(str(item.mRNA)+'\n')


        





